#!/usr/bin/env python
# coding=utf-8
# __author__ = 'Yunchao Ling'

import hashlib
import os
from collections import defaultdict

from Bio import SeqIO
from tqdm import tqdm


def get_md5_value(ss: str):
    my_md5 = hashlib.md5()  # 获取一个MD5的加密算法对象
    my_md5.update(ss.encode("utf-8"))  # 得到MD5消息摘要
    my_md5_Digest = my_md5.hexdigest()  # 以16进制返回消息摘要，32位
    return my_md5_Digest


def check_integrity(gisaid_path: str, ncbi_fasta_path: str):
    ht = defaultdict(list)
    for filename in os.listdir(gisaid_path):
        filepath = os.path.join(gisaid_path, filename)
        if os.path.isfile(filepath) and filename.endswith(".fasta"):
            for seq_record in SeqIO.parse(filepath, "fasta"):
                seq_id = seq_record.description.split("|")[-1]
                fasta = seq_record.seq.upper()
                md5 = get_md5_value(fasta)
                ht[md5].append(seq_id)

    for seq_record in SeqIO.parse(ncbi_fasta_path, "fasta"):
        seq_id = seq_record.id
        fasta = seq_record.seq.upper()
        md5 = get_md5_value(fasta)
        ht[md5].append(seq_id)

    for key in ht.keys():
        print(";".join(ht[key]))


if __name__ == '__main__':
    check_integrity("/Users/genesis/PycharmProjects/ncov/scripts/gisaid_spider/download_data",
                    "/Users/genesis/Downloads/NCBI_Wuhan.fasta")
